package org.example.algorithm.binarytree;

import org.example.algorithm.TreeNode;

import java.util.ArrayList;
import java.util.List;
import java.util.TreeMap;

public class VerticalOrderTraversalSolution2 {

    public List<List<Integer>> verticalTraversal(TreeNode root) {
        List<List<Integer>> res = new ArrayList<>();
        if (root == null) {
            return res;
        }
        TreeMap<Integer, List<int[]>> treeMap = new TreeMap<>();
        dfs(root, 0, 0, treeMap);
        for (List<int[]> list: treeMap.values()) {
            list.sort((a, b)-> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
            List<Integer> newList = new ArrayList<>();
            for (int[] node: list) {
                newList.add(node[1]);
            }
            res.add(newList);
        }
        return res;
    }

    private void dfs(TreeNode root, int row, int column, TreeMap<Integer, List<int[]>> treeMap) {
        treeMap.computeIfAbsent(column, k -> new ArrayList<>()).add(new int[]{row, root.val});
        if (root.left != null) {
            dfs(root.left, row+1, column-1, treeMap);
        }
        if (root.right != null) {
            dfs(root.right, row+1, column+1, treeMap);
        }
    }
}
